# -----------------------------------------------------------------------------
# Copyright (c) 2025, Southeast University (China)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name of the copyright holders nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Authors: Jiajie Xu, Yanfang Zhang, Jiaqi Gao, Leyun Tian
# -----------------------------------------------------------------------------


from Algo_var import *
import networkx as nx
import sys

class Lagrange_Solver():
    """
    TimingGraphWithLagrange uses composition instead of inheritance.
    It wraps around TimingGraph and adds Lagrange multiplier logic.
    """

    
    def __init__(self, contrast=False):
        """
        Constructor to initialize the timing graph with a name.
        
        Args:
            contrast: a contrast cell sizing mode with lambda = 1
        """
        # self.update_lagrange_multipliers(Graph, topo)
        self.contrast = contrast
        self.chosen_cells = {}

    #'''
    #def initialize_lagrange_multipliers(self, Graph, topo):
    def update_lagrange_multipliers(self, Graph, topo):
        """
        Initializes Lagrange multipliers for each node in the graph.

        The method performs the following steps:
        1. Set initial values of λ at sink nodes based on arrival and slack values.
        2. Propagate λ values back through the graph using weighted delay contributions.
        """
        self.Changes_Made = False
        topo_order = topo
        
        # Step 3: Initialize lambda at outputs
        sinks = set()
        initializes = set()


        for node in Graph.nodes:
            if "is_clock_pin" in Graph.nodes[node] and Graph.nodes[node]["is_clock_pin"]:
                Graph.nodes[node]["lambda"] = None
                initializes.add(node)

            # Identify sink nodes (no outgoing edges)
            if Graph.out_degree(node) == 0: 
                arrival_list = Graph.nodes[node].get("arrival", [None, None, None, None]) 
                slack_list = Graph.nodes[node].get("slack", [None, None, None, None])      

                
                # Initialize λ only at flip-flop D pins or primary output ports
                if not Graph.nodes[node].get("is_async_pin", False) or Graph.nodes[node].get("is_port", False):
                    # Choose the more critical timing path (smaller slack)
                    slack_value, arrival_value = 0, 0
                    if slack_list[0] is not None and slack_list[1] is not None:
                        if slack_list[0] < slack_list[1]:
                            slack_value = slack_list[0]
                            arrival_value = arrival_list[0]
                        else:
                            slack_value = slack_list[1]
                            arrival_value = arrival_list[1]

                        # Compute initial λ value
                        if arrival_value + slack_value != 0:
                            lambda_value = arrival_value / (arrival_value + slack_value)
                        else:
                            lambda_value = 0

                        Graph.nodes[node]["lambda"] = lambda_value
                        sinks.add(node)
                    else:
                        Graph.nodes[node]["lambda"] = None
                        initializes.add(node)
                else:
                    Graph.nodes[node]["lambda"] = None
                    initializes.add(node)



        #Step 4: Multiplicative Update
        for u in Graph.nodes:
            for v in Graph[u]:

                edge_dict = Graph.get_edge_data(u, v, default={})

                data = False
                for key, edge_data in Graph[u][v].items():
                    if "is_cell" in edge_data and not edge_data["is_cell"]: #net arc
                        data = True
                        break 

                if data :  
                    continue  

                sense = set()
                max_multiplier = 0.0
                delay_0 = 0.0
                delay_1 = 0.0

                for key, attrs in edge_dict.items():
                    d0, d1 = attrs.get("delay", [None, None, None, None])[:2]

                    if d0 is None and d1 is None:
                            continue
                    
                    delay_0 = max(delay_0, d0 or 0.0)
                    delay_1 = max(delay_1, d1 or 0.0)
                    sense.add(attrs.get("sense_unate",None))


                arrival_i_0,arrival_i_1 = Graph.nodes[u].get('arrival',[None,None,None,None])[:2]
                arrival_j_0,arrival_j_1 = Graph.nodes[v].get('arrival',[None,None,None,None])[:2]

                if arrival_j_0 is None or arrival_j_1 is None :
                    
                    Graph.nodes[u]["lambda"] =  0
                    continue

                if 'positive_unate' in sense :                    
                    multiplier = max(
                        (arrival_i_0 + delay_0)/arrival_j_0,
                        (arrival_i_1 + delay_1)/arrival_j_1
                        )
                    max_multiplier = max(max_multiplier,multiplier)
                if 'negative_unate' in sense :
                    multiplier = max(
                        (arrival_i_0 + delay_1)/arrival_j_1,
                        (arrival_i_1 + delay_0)/arrival_j_0
                        )
                    max_multiplier = max(max_multiplier,multiplier)
                
                
                
                lambda_val = Graph.nodes[u].get("lambda", 1.0)
                if lambda_val is not None:
                    Graph.nodes[u]["lambda"] = lambda_val * max_multiplier

        # Step 5: Backward propagation
        for node in topo_order:
            if node in initializes:
                continue

            node_lambda_value = Graph.nodes[node].get("lambda")
            if node_lambda_value is None:
                continue

           
            total_weight = 0
            pred_weight = []

            # Iterate over predecessor nodes to compute total delay
            for pred in Graph.predecessors(node):
                if pred in initializes:
                    continue

                data = False
                for key, edge_data in Graph[pred][node].items():
                    if "is_cell" in edge_data and not edge_data["is_cell"]: 
                        data = True
                        break  

                if data :  
                    prev_lambda = Graph.nodes[pred].get("lambda", 0)

                    if prev_lambda == 0:
                        Graph.nodes[pred]["lambda"] = node_lambda_value
                    else:
                        Graph.nodes[pred]["lambda"] += node_lambda_value
                    continue



                #total_delay += max_delay
                weight = Graph.nodes[pred].get("lambda",1)
                total_weight += weight
                pred_weight.append((pred, weight))

            #for pred, delay in pred_delays:
            for pred, weight in pred_weight:
                if(total_weight == 0):
                    lambda_contribution = 0
                else:
                    lambda_contribution = node_lambda_value * (weight / total_weight)
                              
                Graph.nodes[pred]["lambda"] = lambda_contribution

        # In contrast mode, all lambda on nodes are set to 1
        if self.contrast:
            for node in Graph.nodes:
                if Graph.nodes[node].get("lambda") != None:
                    Graph.nodes[node]["lambda"] = 1       


        # with open(f"mc_top_new_LM.txt", "w") as f:
        #     for node in Graph.nodes:
        #         lm_value = Graph.nodes[node].get("lambda")
        #         f.write(f"{node}: {lm_value}\n")
    #'''


    def Select_Cell(self, Estimated_subgraph, original_cell, candidate_cells, refname):
        """
        Select the best cell from candidate_cells based on λ-weighted delay sum.

        Args:
            Estimated_subgraph (dict): Key = cell name, Value = nx.MultiDiGraph representing subgraph.
            original_cell (str): Name of the currently used cell.
            candidate_cells (list): List of alternative candidate cells.

        Returns:
            str: The selected cell with the minimum λ × delay sum.
        """
        best_cell = original_cell
        min_cost = float('inf')

        for cell in candidate_cells:
            subgraph = Estimated_subgraph[cell]
            total_cost = 0.0

            for u, v in subgraph.edges():
                edge_dict = subgraph.get_edge_data(u, v, default={})
                
                attrs = edge_dict.get(0, {})
                if not attrs.get("is_cell", False):
                    continue
                
                #max delay
                max_cell_delay_d0 = 0.0
                max_cell_delay_d1 = 0.0
                for key, attrs in edge_dict.items():
                    d0, d1 = attrs.get("delay", [0.0, 0.0, 0.0, 0.0])[:2]
                    max_cell_delay_d0 = max(max_cell_delay_d0, d0 or 0.0)
                    max_cell_delay_d1 = max(max_cell_delay_d1, d1 or 0.0)

                delay_net_d0 = 0.0
                delay_net_d1 = 0.0
                for _, next_node, edge_data in subgraph.out_edges(v, data=True):
                    if not edge_data.get("is_cell", True):  
                        d0_n, d1_n = edge_data.get("delay", [0.0, 0.0, 0.0, 0.0])[:2]
                        delay_net_d0 = max(delay_net_d0, d0_n or 0.0)
                        delay_net_d1 = max(delay_net_d1, d1_n or 0.0)

                max_delay=max(max_cell_delay_d0+delay_net_d0,max_cell_delay_d1+delay_net_d1)

                
                lambda_val = subgraph.nodes[u].get("lambda", 0.0)
                if lambda_val is None:
                    lambda_val = 0.0

                total_cost += lambda_val * max_delay

            if total_cost < min_cost:
                min_cost = total_cost
                best_cell = cell

        self.chosen_cells[refname] = best_cell
        
        if best_cell != original_cell:
            self.Changes_Made = True
        
        return best_cell
    
    def Output_ECO_Change_List(self, filename = "ECO_ChangeList"):
        """
        Output ECO Changelist for PrimeTime.
        """
        with open(filename, "w") as outfile:
            outfile.write("current_instance\n")
            for key, value in self.chosen_cells.items():
                outfile.write(f"size_cell {key} {value}\n")


# Example usage
#if __name__ == "__main__":
    #TG = TimingGraphWithLagrange("mc_top")
    #TG = TimingGraphWithLagrange("sasc_top")
    #TG = TimingGraphWithLagrange("mc_top")

    #TG.initialize_lagrange_multipliers()
